# 文件名: 卡尺ui.py (V9.1 - RANSAC适配版)
import sys
import cv2
import numpy as np
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QFileDialog, QTabWidget, QSlider, QComboBox,
    QGroupBox, QFormLayout, QSpinBox, QCheckBox, QScrollArea
)
from PySide6.QtGui import QPixmap, QImage, QPainter, QPen, QColor, QFont
from PySide6.QtCore import Qt, QPoint, QRect, Signal, QPointF

# 导入V9版本的函数
from opencv卡尺工具 import line_edge_detection, caliper_circle_detection

def convert_cv_qt(cv_img):
    if cv_img is None: return QPixmap()
    rgb_image = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
    h, w, ch = rgb_image.shape
    bytes_per_line = ch * w
    convert_to_Qt_format = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
    return QPixmap.fromImage(convert_to_Qt_format)

class InteractiveImageLabel(QLabel):
    lineRoiDefined = Signal(QPoint, QPoint)
    circleRoiDefined = Signal(QPoint, int)
    def __init__(self, scroll_area, parent=None):
        super().__init__(parent)
        self.scroll_area = scroll_area
        self.setMouseTracking(True)
        self.drawing = False; self.start_point_img = QPoint(); self.end_point_img = QPoint(); self.draw_mode = 'none'
        self.original_pixmap = None; self.scale_factor = 1.0
        self.user_line_roi = None; self.user_circle_roi = None
        # V9.1 适配：新增 inlier_points 用于区分内外点
        self.detected_edge_points = []; self.detected_inlier_points = []
        self.detected_line = None; self.detected_circle = None
        
    def set_pixmap(self, pixmap):
        self.original_pixmap = pixmap; self.scale_factor = 1.0
        self.clear_all_overlays(); self.update_scaled_image()
        self.scroll_area.horizontalScrollBar().setValue(0); self.scroll_area.verticalScrollBar().setValue(0)
        
    def update_pixmap(self, pixmap):
        self.original_pixmap = pixmap; self.update_scaled_image()
        
    def set_user_line_roi(self, pt1, pt2):
        self.user_line_roi = (QPoint(pt1[0], pt1[1]), QPoint(pt2[0], pt2[1])); self.user_circle_roi = None; self.clear_detection_results()
        
    def set_user_circle_roi(self, center, radius):
        self.user_circle_roi = (QPoint(center[0], center[1]), radius); self.user_line_roi = None; self.clear_detection_results()

    # V9.1 适配：修改函数签名以接收 inlier_points
    def set_detection_results(self, all_points, inlier_points, detected_geom):
        self.detected_edge_points = all_points
        self.detected_inlier_points = inlier_points # 存储内点
        if self.user_line_roi: self.detected_line = detected_geom; self.detected_circle = None
        elif self.user_circle_roi: self.detected_circle = detected_geom; self.detected_line = None
        self.update()

    def clear_detection_results(self):
        self.detected_edge_points = []; self.detected_inlier_points = [] # 同时清空内点
        self.detected_line = None; self.detected_circle = None; self.update()
        
    def clear_all_overlays(self):
        self.user_line_roi = None; self.user_circle_roi = None; self.clear_detection_results()
        
    def wheelEvent(self, event):
        if self.original_pixmap is None or self.original_pixmap.isNull(): return
        zoom_in_factor = 1.15; zoom_out_factor = 1 / zoom_in_factor; old_scale_factor = self.scale_factor
        if event.angleDelta().y() > 0: self.scale_factor *= zoom_in_factor
        else: self.scale_factor *= zoom_out_factor
        self.scale_factor = max(0.1, min(self.scale_factor, 10.0))
        if abs(self.scale_factor - old_scale_factor) < 1e-5: return
        mouse_pos_viewport = event.position(); h_bar = self.scroll_area.horizontalScrollBar(); v_bar = self.scroll_area.verticalScrollBar()
        h_ratio = (mouse_pos_viewport.x() + h_bar.value()) / self.width() if self.width() > 0 else 0
        v_ratio = (mouse_pos_viewport.y() + v_bar.value()) / self.height() if self.height() > 0 else 0
        self.update_scaled_image()
        new_h_val = h_ratio * self.width() - mouse_pos_viewport.x(); new_v_val = v_ratio * self.height() - mouse_pos_viewport.y()
        h_bar.setValue(int(new_h_val)); v_bar.setValue(int(new_v_val))
        
    def paintEvent(self, event):
        super().paintEvent(event);
        if self.original_pixmap is None: return
        painter = QPainter(self); painter.setRenderHint(QPainter.Antialiasing)
        pen_roi = QPen(QColor(0, 255, 0), 2, Qt.SolidLine); painter.setPen(pen_roi)
        if self.user_line_roi:
            pt1_label = self.map_image_to_pos(self.user_line_roi[0]); pt2_label = self.map_image_to_pos(self.user_line_roi[1])
            if pt1_label and pt2_label: painter.drawLine(pt1_label, pt2_label)
        if self.user_circle_roi:
            center_label = self.map_image_to_pos(self.user_circle_roi[0]); radius_label = self.user_circle_roi[1] * self.scale_factor
            if center_label: painter.drawEllipse(center_label, radius_label, radius_label)
        
        # V9.1 适配：区分绘制内点和外点
        pen_inlier = QPen(QColor(0, 255, 255), 1)   # 内点: 黄色/青色
        pen_outlier = QPen(QColor(255, 0, 255), 1)  # 外点: 紫色
        marker_size = 5
        
        inlier_set = {tuple(map(int, pt)) for pt in self.detected_inlier_points}

        for pt_tuple in self.detected_edge_points:
            pt_int_tuple = tuple(map(int, pt_tuple))
            p = self.map_image_to_pos(QPointF(pt_tuple[0], pt_tuple[1]))
            if p:
                if pt_int_tuple in inlier_set:
                    painter.setPen(pen_inlier) # 设置为内点颜色
                else:
                    painter.setPen(pen_outlier) # 设置为外点颜色
                
                # 绘制十字标记
                painter.drawLine(p.x() - marker_size, p.y(), p.x() + marker_size, p.y())
                painter.drawLine(p.x(), p.y() - marker_size, p.x(), p.y() + marker_size)

        pen_result = QPen(QColor(255, 0, 0), 2); painter.setPen(pen_result)
        if self.detected_line:
            pt1_label = self.map_image_to_pos(QPoint(*self.detected_line[0])); pt2_label = self.map_image_to_pos(QPoint(*self.detected_line[1]))
            if pt1_label and pt2_label: painter.drawLine(pt1_label, pt2_label)
        if self.detected_circle:
            center, radius = self.detected_circle; center_label = self.map_image_to_pos(QPointF(center[0], center[1])); radius_label = radius * self.scale_factor
            if center_label: painter.drawEllipse(center_label, radius_label, radius_label)
            
        if self.drawing and not self.start_point_img.isNull():
            pen_draw = QPen(Qt.cyan, 2, Qt.DashLine); painter.setPen(pen_draw)
            start_pos_label = self.map_image_to_pos(self.start_point_img); end_pos_label = self.map_image_to_pos(self.end_point_img)
            if start_pos_label and end_pos_label:
                if self.draw_mode == 'line': painter.drawLine(start_pos_label, end_pos_label)
                elif self.draw_mode == 'circle':
                    radius_label = int(np.linalg.norm(np.array([start_pos_label.x(), start_pos_label.y()]) - np.array([end_pos_label.x(), end_pos_label.y()])))
                    painter.drawEllipse(start_pos_label, radius_label, radius_label)
        painter.end()
        
    def set_draw_mode(self, mode): self.draw_mode = mode; self.start_point_img = QPoint(); self.end_point_img = QPoint(); self.update()
    
    def update_scaled_image(self):
        if self.original_pixmap and not self.original_pixmap.isNull():
            new_width = int(self.original_pixmap.width() * self.scale_factor); new_height = int(self.original_pixmap.height() * self.scale_factor)
            scaled_pixmap = self.original_pixmap.scaled(new_width, new_height, Qt.KeepAspectRatio, Qt.FastTransformation)
            super().setPixmap(scaled_pixmap); self.resize(scaled_pixmap.size())
            
    def map_pos_to_image(self, pos):
        if self.pixmap() is None or self.pixmap().isNull(): return None
        pixmap_size = self.pixmap().size(); label_size = self.size()
        offset_x = (label_size.width() - pixmap_size.width()) / 2; offset_y = (label_size.height() - pixmap_size.height()) / 2
        img_scaled_x = pos.x() - offset_x; img_scaled_y = pos.y() - offset_y
        original_x = img_scaled_x / self.scale_factor; original_y = img_scaled_y / self.scale_factor
        if self.original_pixmap and 0 <= original_x < self.original_pixmap.width() and 0 <= original_y < self.original_pixmap.height(): return QPoint(int(original_x), int(original_y))
        return None
        
    def map_image_to_pos(self, img_pos):
        if self.pixmap() is None or self.pixmap().isNull(): return None
        scaled_x = img_pos.x() * self.scale_factor; scaled_y = img_pos.y() * self.scale_factor
        pixmap_size = self.pixmap().size(); label_size = self.size()
        offset_x = (label_size.width() - pixmap_size.width()) / 2; offset_y = (label_size.height() - pixmap_size.height()) / 2
        return QPointF(scaled_x + offset_x, scaled_y + offset_y)
        
    def mousePressEvent(self, event):
        pos_in_image = self.map_pos_to_image(event.pos())
        if event.button() == Qt.LeftButton and self.draw_mode != 'none' and pos_in_image: self.drawing = True; self.start_point_img = pos_in_image; self.end_point_img = pos_in_image; self.update()
        
    def mouseMoveEvent(self, event):
        pos_in_image = self.map_pos_to_image(event.pos())
        if self.drawing and pos_in_image: self.end_point_img = pos_in_image; self.update()
        
    def mouseReleaseEvent(self, event):
        pos_in_image = self.map_pos_to_image(event.pos())
        if event.button() == Qt.LeftButton and self.drawing and pos_in_image:
            self.drawing = False; self.end_point_img = pos_in_image
            if self.draw_mode == 'line': self.lineRoiDefined.emit(self.start_point_img, self.end_point_img)
            elif self.draw_mode == 'circle':
                radius = int(np.linalg.norm(np.array([self.start_point_img.x(), self.start_point_img.y()]) - np.array([self.end_point_img.x(), self.end_point_img.y()])))
                self.circleRoiDefined.emit(self.start_point_img, radius)
            self.update()

class CaliperApp(QMainWindow):
    def __init__(self):
        super().__init__()
        # V9.1 适配：更新窗口标题
        self.setWindowTitle("OpenCV 卡尺工具 (PySide6 UI) - V9 RANSAC版")
        self.setGeometry(100, 100, 1200, 800)
        self.cv_image = None; self.line_pt1 = None; self.line_pt2 = None; self.circle_center = None; self.circle_radius = None
        main_widget = QWidget(); self.setCentralWidget(main_widget); main_layout = QHBoxLayout(main_widget)
        self.create_control_panel(); main_layout.addWidget(self.control_panel, 1)
        self.image_display_area = QScrollArea(); self.image_display_area.setWidgetResizable(True)
        self.image_label = InteractiveImageLabel(self.image_display_area)
        self.image_label.setAlignment(Qt.AlignCenter); self.image_display_area.setWidget(self.image_label)
        main_layout.addWidget(self.image_display_area, 3)
        self.connect_signals(); self.update_ui_state()
        
    def create_control_panel(self):
        self.control_panel = QWidget(); controls_layout = QVBoxLayout(self.control_panel); file_group = QGroupBox("文件"); file_layout = QVBoxLayout()
        self.btn_load_image = QPushButton("加载图像"); file_layout.addWidget(self.btn_load_image); file_group.setLayout(file_layout); controls_layout.addWidget(file_group)
        self.tabs = QTabWidget(); self.create_line_tab(); self.create_circle_tab(); self.tabs.addTab(self.line_tab, "直线检测"); self.tabs.addTab(self.circle_tab, "圆形检测")
        controls_layout.addWidget(self.tabs); controls_layout.addStretch()
    
    def create_line_tab(self):
        self.line_tab = QWidget(); layout = QVBoxLayout(self.line_tab); form_layout = QFormLayout(); self.btn_define_line_roi = QPushButton("定义引导线 ROI")
        self.line_caliper_count = IntSliderSpinBox(min_val=1, max_val=100, default_val=15)
        self.line_search_width = IntSliderSpinBox(min_val=10, max_val=200, default_val=80)
        self.line_projection_length = IntSliderSpinBox(min_val=5, max_val=100, default_val=30)
        self.line_polarity = QComboBox(); self.line_polarity.addItems(['positive', 'negative', 'both']); self.line_polarity.setCurrentText('negative')
        self.line_show_roi = QCheckBox("显示检测ROI"); self.line_show_roi.setChecked(True); form_layout.addRow(self.btn_define_line_roi)
        form_layout.addRow("卡尺数量:", self.line_caliper_count)
        form_layout.addRow("搜索宽度:", self.line_search_width)
        form_layout.addRow("投影长度:", self.line_projection_length)
        form_layout.addRow("边缘极性:", self.line_polarity); form_layout.addRow(self.line_show_roi)
        self.btn_run_line = QPushButton("运行直线检测"); layout.addLayout(form_layout); layout.addWidget(self.btn_run_line); layout.addStretch()

    def create_circle_tab(self):
        self.circle_tab = QWidget(); layout = QVBoxLayout(self.circle_tab); form_layout = QFormLayout(); self.btn_define_circle_roi = QPushButton("定义引导圆 ROI")
        self.circle_caliper_count = IntSliderSpinBox(min_val=4, max_val=100, default_val=16)
        self.circle_search_width = IntSliderSpinBox(min_val=10, max_val=200, default_val=40)
        self.circle_projection_length = IntSliderSpinBox(min_val=5, max_val=100, default_val=15)
        self.circle_polarity = QComboBox(); self.circle_polarity.addItems(['positive', 'negative', 'both']); self.circle_polarity.setCurrentText('positive')
        self.circle_show_roi = QCheckBox("显示检测ROI"); self.circle_show_roi.setChecked(True)
        # V9.1 适配：添加RANSAC参数的UI控件
        self.ransac_threshold = IntSliderSpinBox(min_val=1, max_val=20, default_val=2)
        self.ransac_iterations = IntSliderSpinBox(min_val=50, max_val=500, default_val=100)
        
        form_layout.addRow(self.btn_define_circle_roi)
        form_layout.addRow("卡尺数量:", self.circle_caliper_count)
        form_layout.addRow("搜索宽度:", self.circle_search_width)
        form_layout.addRow("投影长度:", self.circle_projection_length)
        form_layout.addRow("边缘极性:", self.circle_polarity)
        form_layout.addRow("RANSAC阈值:", self.ransac_threshold)
        form_layout.addRow("RANSAC迭代:", self.ransac_iterations)
        form_layout.addRow(self.circle_show_roi)
        self.btn_run_circle = QPushButton("运行圆形检测"); layout.addLayout(form_layout); layout.addWidget(self.btn_run_circle); layout.addStretch()
        
    def connect_signals(self):
        self.btn_load_image.clicked.connect(self.load_image); self.btn_define_line_roi.clicked.connect(self.define_line_roi); self.btn_run_line.clicked.connect(self.run_line_detection)
        self.btn_define_circle_roi.clicked.connect(self.define_circle_roi); self.btn_run_circle.clicked.connect(self.run_circle_detection)
        self.image_label.lineRoiDefined.connect(self.on_line_roi_defined); self.image_label.circleRoiDefined.connect(self.on_circle_roi_defined)
        
    def update_ui_state(self):
        is_image_loaded = self.cv_image is not None
        self.btn_define_line_roi.setEnabled(is_image_loaded); self.btn_define_circle_roi.setEnabled(is_image_loaded)
        self.btn_run_line.setEnabled(is_image_loaded and self.line_pt1 is not None)
        self.btn_run_circle.setEnabled(is_image_loaded and self.circle_center is not None)
        
    def load_image(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "选择图像文件", "", "Image Files (*.png *.jpg *.jpeg *.bmp)")
        if file_name:
            self.cv_image = cv2.imdecode(np.fromfile(file_name, dtype=np.uint8), cv2.IMREAD_COLOR)
            if self.cv_image is not None:
                self.image_label.set_pixmap(convert_cv_qt(self.cv_image)); self.statusBar().showMessage(f"图像加载成功: {file_name}", 5000)
                self.line_pt1 = self.line_pt2 = self.circle_center = self.circle_radius = None
            else: self.statusBar().showMessage("错误: 无法加载图像", 5000)
        self.update_ui_state()
        
    def define_line_roi(self): self.image_label.set_draw_mode('line'); self.statusBar().showMessage("请在图像上拖拽以定义引导线ROI", 5000)
    
    def on_line_roi_defined(self, pt1, pt2):
        self.line_pt1 = (pt1.x(), pt1.y()); self.line_pt2 = (pt2.x(), pt2.y()); self.statusBar().showMessage(f"引导线ROI定义完成", 5000)
        self.image_label.set_draw_mode('none'); self.image_label.set_user_line_roi(self.line_pt1, self.line_pt2); self.update_ui_state()
        
    def define_circle_roi(self): self.image_label.set_draw_mode('circle'); self.statusBar().showMessage("请在图像上拖拽以定义引导圆ROI", 5000)
    
    def on_circle_roi_defined(self, center, radius):
        self.circle_center = (center.x(), center.y()); self.circle_radius = radius; self.statusBar().showMessage(f"引导圆ROI定义完成", 5000)
        self.image_label.set_draw_mode('none'); self.image_label.set_user_circle_roi(self.circle_center, self.circle_radius); self.update_ui_state()
        
    def run_line_detection(self):
        if self.cv_image is None or self.line_pt1 is None: return
        params = {
            "caliper_count": self.line_caliper_count.value(),
            "search_width": self.line_search_width.value(),
            "projection_length": self.line_projection_length.value(),
            "edge_polarity": self.line_polarity.currentText(),
            "show_roi": self.line_show_roi.isChecked()
        }
        self.statusBar().showMessage("正在执行直线检测..."); QApplication.processEvents()
        detected_line, edge_points, vis_img = line_edge_detection(self.cv_image.copy(), self.line_pt1, self.line_pt2, **params)
        self.image_label.update_pixmap(convert_cv_qt(vis_img))
        if detected_line:
            self.statusBar().showMessage("直线检测成功！", 5000)
            # V9.1 适配：由于直线检测没有内外点之分，将所有点都作为内点传入
            self.image_label.set_detection_results(edge_points, edge_points, detected_line)
        else:
            self.statusBar().showMessage("直线检测失败", 5000); self.image_label.clear_detection_results()
            
    def run_circle_detection(self):
        if self.cv_image is None or self.circle_center is None: return
        params = {
            "caliper_count": self.circle_caliper_count.value(),
            "search_width": self.circle_search_width.value(),
            "projection_length": self.circle_projection_length.value(),
            "edge_polarity": self.circle_polarity.currentText(),
            "show_roi": self.circle_show_roi.isChecked(),
            # V9.1 适配：从UI获取RANSAC参数
            "ransac_threshold": self.ransac_threshold.value(),
            "ransac_iterations": self.ransac_iterations.value()
        }
        self.statusBar().showMessage("正在执行圆形检测..."); QApplication.processEvents()
        
        # V9.1 适配：用4个变量接收函数返回的4个值
        (final_center, final_radius), all_edge_points, inlier_points, vis_img = caliper_circle_detection(
            self.cv_image.copy(), self.circle_center, self.circle_radius, **params)
            
        self.image_label.update_pixmap(convert_cv_qt(vis_img))
        if final_center is not None and final_radius is not None:
            message = f"圆形检测成功！找到 {len(all_edge_points)} 个点, 其中 {len(inlier_points)} 个是内点"
            self.statusBar().showMessage(message, 5000)
            # V9.1 适配：传入所有点、内点和几何结果
            self.image_label.set_detection_results(all_edge_points, inlier_points, (final_center, final_radius))
        else:
            self.statusBar().showMessage("圆形检测失败", 5000); self.image_label.clear_detection_results()

class IntSliderSpinBox(QWidget):
    # ... [这部分代码无需修改] ...
    valueChanged = Signal(int)
    def __init__(self, orientation=Qt.Horizontal, min_val=0, max_val=100, default_val=50, is_odd=False):
        super().__init__(); self.is_odd = is_odd; layout = QHBoxLayout(self); layout.setContentsMargins(0, 0, 0, 0)
        self.slider = QSlider(orientation); self.slider.setRange(min_val, max_val); self.spinbox = QSpinBox(); self.spinbox.setRange(min_val, max_val)
        if is_odd: self.slider.setSingleStep(2); self.spinbox.setSingleStep(2); default_val = max(min_val, default_val - 1) if default_val % 2 == 0 else default_val
        self.slider.setValue(default_val); self.spinbox.setValue(default_val)
        self.slider.valueChanged.connect(self.on_slider_value_changed); self.spinbox.valueChanged.connect(self.on_spinbox_value_changed)
        layout.addWidget(self.slider); layout.addWidget(self.spinbox)
    def on_slider_value_changed(self, value):
        if self.is_odd and value % 2 == 0: value = max(self.slider.minimum(), value - 1); self.slider.setValue(value); return
        self.spinbox.setValue(value); self.valueChanged.emit(value)
    def on_spinbox_value_changed(self, value): self.slider.setValue(value)
    def value(self): return self.spinbox.value()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CaliperApp()
    window.show()
    sys.exit(app.exec())
